from enum import Enum
from mlagents_envs.exception import UnityCommunicationException
from mlagents_envs.side_channel.side_channel import (
    SideChannel,
    IncomingMessage,
    OutgoingMessage,
)
import uuid

class EnvSettingType(Enum):
    SET_PARAMETER = 1
    HARD_RESET = 2
    OTHER = 3

class EnvParaType(Enum):
    BOOL = 1
    INT = 2
    FLOAT = 3
    LIST_FLOAT = 4
    RANDOM_NORMAL = 5
    RANDOM_UNIFORM = 6

# *** Source definition of the environment parameters ***
class EnvPara(Enum):
    seed = "seed"
    playground_number = "playgroundNumber"
    target_pos_x = "targetPosX"
    target_pos_y = "targetPosY"
    target_pos_z = "targetPosZ"
    spawn_pos_x = "spawnPosX"
    spawn_pos_y = "spawnPosY"
    spawn_pos_z = "spawnPosZ"
    enable_random_airplane_spawn_rotation = "enableRandomAirplaneSpawnRotation"
    enable_boundary_check = "enableBoundaryCheck"
    boundary_forward = "boundaryForward"
    boundary_backward = "boundaryBackward"
    boundary_up = "boundaryUp"
    boundary_down = "boundaryDown"
    boundary_left = "boundaryLeft"
    boundary_right = "boundaryRight"
    max_distance_to_target = "maxDistanceToTarget"
    agent_number_per_playground = "agentNumberPerPlayground"
    badspheres_pos_z = "badSpheresPosZ"
    badsphere_pos_x = "badSpherePosX"
    badsphere_pos_y = "badSpherePosY"
    badsphere_scale = "badSphereScale"
    distance_reward_scale = "distanceRewardScale"
    angular_velocity_punish_scale = "angularVelocityPunishScale"
    time_consume_punish_scale = "timeConsumePunishScale"
    unified_episode = "unifiedEpisode"
    is_crash_end_episode = "isCrashEndEpisode"


class EnvironmentSettingChannel(SideChannel):

    def __init__(self) -> None:
        super().__init__(uuid.UUID("621f0a70-4f87-11ea-a6bf-784f4387d1f8"))

    def on_message_received(self, msg: IncomingMessage) -> None:
        raise UnityCommunicationException(
            "The EnvironmentSettingChannel received a message from Unity, "
            + "this should not have happend."
        )
    
    def set_environment_parameter(self, key: EnvPara, env_para_type: EnvParaType, value: object):
        msg = OutgoingMessage()
        msg.write_int32(EnvSettingType.SET_PARAMETER.value)
        msg.write_string(key.value)
        msg.write_int32(env_para_type.value)
        if(env_para_type == EnvParaType.BOOL):
            assert isinstance(value, bool), "Parameter type mismatch."
            msg.write_bool(value)
        elif(env_para_type == EnvParaType.INT):
            assert isinstance(value, int), "Parameter type mismatch."
            msg.write_int32(value)
        elif(env_para_type == EnvParaType.FLOAT):
            assert isinstance(value, float), "Parameter type mismatch."
            msg.write_float32(value)
        elif(env_para_type == EnvParaType.LIST_FLOAT):
            assert isinstance(value, list), "Parameter type mismatch."
            msg.write_float32_list(value)
        elif(env_para_type == EnvParaType.RANDOM_NORMAL or env_para_type == EnvParaType.RANDOM_UNIFORM):
            assert isinstance(value, list) and len(value)==3, "Parameter type mismatch."
            msg.write_float32_list(value)

        super().queue_message_to_send(msg)

    # 2021.4.10 放弃hardset功能的实现
    def do_hard_reset(self):
        msg = OutgoingMessage()
        msg.write_int32(EnvSettingType.HARD_RESET.value)
        super().queue_message_to_send(msg)